/*
  KeePass Password Safe - The Open-Source Password Manager
  Copyright (C) 2003-2011 Dominik Reichl <dominik.reichl@t-online.de>

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "StdAfx.h"


#include <map>
#include <set>

#include "PwUtil.h"
#include "MemUtil.h"
#include "StrUtil.h"

#include "PopularPasswords.h"
#include "../../Common/Crypto/ARCFour.h"
#include <math.h>


#define CHARSPACE_ESCAPE      60
#define CHARSPACE_ALPHA       26
#define CHARSPACE_NUMBER      10
#define CHARSPACE_SIMPSPECIAL 16
#define CHARSPACE_EXTSPECIAL  17
#define CHARSPACE_HIGH       112



static const BYTE g_uuidZero[16] = { 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0 };

CPwUtil::CPwUtil()
{
}

// Very simple password quality estimation function
DWORD CPwUtil::EstimatePasswordBits(LPCTSTR lpPassword)
{
	BOOL bChLower = FALSE, bChUpper = FALSE, bChNumber = FALSE;
	BOOL bChSimpleSpecial = FALSE, bChExtSpecial = FALSE, bChHigh = FALSE;
	BOOL bChEscape = FALSE;
	double dblEffectiveLength = 0.0;
	std::map<TCHAR, unsigned int> vCharCounts;
	std::map<int, unsigned int> vDifferences;

	ATLASSERT(lpPassword != NULL); if(lpPassword == NULL) return 0;

	const DWORD dwLen = static_cast<DWORD>(_tcslen(lpPassword));
	if(dwLen == 0) return 0; // Zero bits of information :)

	for(DWORD i = 0; i < dwLen; ++i) // Get character types
	{
		const TCHAR tch = lpPassword[i];

		if((tch > 0) && (tch < _T(' '))) bChEscape = TRUE;
		if((tch >= _T('A')) && (tch <= _T('Z'))) bChUpper = TRUE;
		if((tch >= _T('a')) && (tch <= _T('z'))) bChLower = TRUE;
		if((tch >= _T('0')) && (tch <= _T('9'))) bChNumber = TRUE;
		if((tch >= _T(' ')) && (tch <= _T('/'))) bChSimpleSpecial = TRUE;
		if((tch >= _T(':')) && (tch <= _T('@'))) bChExtSpecial = TRUE;
		if((tch >= _T('[')) && (tch <= _T('`'))) bChExtSpecial = TRUE;
		if((tch >= _T('{')) && (tch <= _T('~'))) bChExtSpecial = TRUE;
		if((tch <= 0) || (tch > _T('~'))) bChHigh = TRUE;

		double dblDiffFactor = 1.0;
		if(i >= 1)
		{
			const int iDiff = (int)tch - (int)lpPassword[i - 1];

			if(vDifferences.find(iDiff) == vDifferences.end())
				vDifferences[iDiff] = 1;
			else
			{
				vDifferences[iDiff] = vDifferences[iDiff] + 1;
				dblDiffFactor /= (double)vDifferences[iDiff];
			}
		}

		if(vCharCounts.find(tch) == vCharCounts.end())
		{
			vCharCounts[tch] = 1;
			dblEffectiveLength += dblDiffFactor;
		}
		else
		{
			vCharCounts[tch] = vCharCounts[tch] + 1;
			dblEffectiveLength += dblDiffFactor * (1.0 / (double)vCharCounts[tch]);
		}
	}

	DWORD dwCharSpace = 0;
	if(bChEscape == TRUE) dwCharSpace += CHARSPACE_ESCAPE;
	if(bChUpper == TRUE) dwCharSpace += CHARSPACE_ALPHA;
	if(bChLower == TRUE) dwCharSpace += CHARSPACE_ALPHA;
	if(bChNumber == TRUE) dwCharSpace += CHARSPACE_NUMBER;
	if(bChSimpleSpecial == TRUE) dwCharSpace += CHARSPACE_SIMPSPECIAL;
	if(bChExtSpecial == TRUE) dwCharSpace += CHARSPACE_EXTSPECIAL;
	if(bChHigh == TRUE) dwCharSpace += CHARSPACE_HIGH;

	ATLASSERT(dwCharSpace != 0); if(dwCharSpace == 0) return 0;

	const double dblBitsPerChar = log((double)dwCharSpace) / log(2.00);
	double dblRating = dblBitsPerChar * dblEffectiveLength;

	if(IsPopularPassword(lpPassword)) dblRating /= 8.0;

	const DWORD dwBits = static_cast<DWORD>(ceil(dblRating));

	ATLASSERT(dwBits != 0);
	return dwBits;
}

BOOL CPwUtil::LoadHexKey32(FILE *fp, BYTE *pBuf)
{
	char buf[65], ch1, ch2;
	BYTE bt;

	ATLASSERT(fp != NULL); if(fp == NULL) return FALSE;
	ATLASSERT(pBuf != NULL); if(pBuf == NULL) return FALSE;

	buf[64] = 0;
	if(fread(buf, 1, 64, fp) != 64) { ATLASSERT(FALSE); return FALSE; }

	for(int i = 0; i < 32; ++i)
	{
		ch1 = buf[i * 2];
		ch2 = buf[i * 2 + 1];

		if(CPwUtil::ConvertStrToHex(ch1, ch2, bt) == FALSE) return FALSE;

		pBuf[i] = bt;
	}

	mem_erase((BYTE *)buf, 64);
	return TRUE;
}

BOOL CPwUtil::SaveHexKey32(FILE *fp, BYTE *pBuf)
{
	char buf[65], ch1, ch2;
	BYTE bt;

	ATLASSERT(fp != NULL); if(fp == NULL) return FALSE;
	ATLASSERT(pBuf != NULL); if(pBuf == NULL) return FALSE;

	buf[64] = 0;

	for(int i = 0; i < 32; ++i)
	{
		bt = pBuf[i];

		CPwUtil::ConvertHexToStr(bt, ch1, ch2);

		buf[i * 2] = ch1;
		buf[i * 2 + 1] = ch2;
	}

	fwrite(buf, 1, 64, fp);

	mem_erase((BYTE *)buf, 64);
	return TRUE;
}

//LPCTSTR CPwUtil::GetUniCvtPtr(LPCTSTR lpBase, BOOL bTranslate)
//{
//	if(lpBase == NULL) return PWU_CVT_EX; // NULL is allowed, return unique default pointer
//
//	if(bTranslate != FALSE) return _TRL(lpBase);
//	return lpBase; // Untranslated
//}

BOOL CPwUtil::ConvertStrToHex(char ch1, char ch2, BYTE& bt)
{
	if((ch1 >= '0') && (ch1 <= '9')) bt = static_cast<BYTE>(ch1 - '0');
	else if((ch1 >= 'a') && (ch1 <= 'f')) bt = static_cast<BYTE>(ch1 - 'a' + 10);
	else if((ch1 >= 'A') && (ch1 <= 'F')) bt = static_cast<BYTE>(ch1 - 'A' + 10);
	else return FALSE;

	bt <<= 4;

	if((ch2 >= '0') && (ch2 <= '9')) bt |= static_cast<BYTE>(ch2 - '0');
	else if((ch2 >= 'a') && (ch2 <= 'f')) bt |= static_cast<BYTE>(ch2 - 'a' + 10);
	else if((ch2 >= 'A') && (ch2 <= 'F')) bt |= static_cast<BYTE>(ch2 - 'A' + 10);
	else return FALSE;

	return TRUE;
}

void CPwUtil::ConvertHexToStr(BYTE bt, char& ch1, char& ch2)
{
	char chq = static_cast<char>(bt >> 4);
	if(chq < 10) ch1 = static_cast<char>(chq + '0');
	else ch1 = static_cast<char>(chq - 10 + 'a');

	chq = static_cast<char>(bt & 0x0F);
	if(chq < 10) ch2 = static_cast<char>(chq + '0');
	else ch2 = static_cast<char>(chq - 10 + 'a');
}

CString CPwUtil::FormatError(int nErrorCode, DWORD dwFlags)
{
	CString str;

	if(nErrorCode != PWE_SUCCESS)
	{
		TCHAR tszTemp[24];
		_stprintf_s(tszTemp, _countof(tszTemp), _T("%08X"),
			static_cast<unsigned int>(nErrorCode));

		if((dwFlags & PWFF_NO_INTRO) == 0)
		{
			str = BkString::Get(91); str += _T("!\r\n");
		}

		str +=BkString::Get(92); str += _T(": 0x");
		str += tszTemp;
		str += _T(".");

		if((dwFlags & PWFF_NO_INTRO) == 0) str += _T("\r\n\r\n");
		else str += _T("\r\n");
	}

	if(((nErrorCode == PWE_INVALID_KEY) && ((dwFlags & PWFF_INVKEY_WITH_CODE) == 0)) ||
		(nErrorCode == PWE_UNSUPPORTED_KDBX) || ((dwFlags & PWFF_MAIN_TEXT_ONLY) != 0) ||
		(nErrorCode == PWE_GETLASTERROR))
		str.Empty();

	switch(nErrorCode)
	{
	case PWE_UNKNOWN:
		str += BkString::Get(93);
		break;
	case PWE_SUCCESS:
		str += BkString::Get(94);
		break;
	case PWE_INVALID_PARAM:
		str += BkString::Get(95);
		break;
	case PWE_NO_MEM:
		str += BkString::Get(96);
		break;
	case PWE_INVALID_KEY:
		str += BkString::Get(97);
		break;
	case PWE_NOFILEACCESS_READ:
		str += BkString::Get(98);
		break;
	case PWE_NOFILEACCESS_WRITE:
		str += BkString::Get(99);
		break;
	case PWE_FILEERROR_READ:
		str +=BkString::Get(500);
		break;
	case PWE_FILEERROR_WRITE:
		str += BkString::Get(501);
		break;
	case PWE_INVALID_RANDOMSOURCE:
		str += BkString::Get(502); str += _T(": ");
		str += BkString::Get(503);
		break;
	case PWE_INVALID_FILESTRUCTURE:
		str += BkString::Get(504);
		break;
	case PWE_CRYPT_ERROR:
		str += BkString::Get(505);
		break;
	case PWE_INVALID_FILESIZE:
		str += BkString::Get(504);
		break;
	case PWE_INVALID_FILESIGNATURE:
		str += BkString::Get(504);
		break;
	case PWE_INVALID_FILEHEADER:
		str +=BkString::Get(504);
		break;
	case PWE_NOFILEACCESS_READ_KEY:
		str += BkString::Get(98);
		break;
	case PWE_KEYPROV_INVALID_KEY:
		str += _T("The key provider plugin did not supply a valid key");
		break;
	case PWE_FILEERROR_VERIFY:
		str += BkString::Get(501); // Same as write
		break;
	case PWE_UNSUPPORTED_KDBX:
		str += _T("KeePass 1.x cannot open KDBX files.");
		str += _T("\r\n\r\n");
		str += _T("Use the KeePass 2.x 'Export' feature to migrate this file");
		break;
	case PWE_GETLASTERROR:
		str += CPwUtil::FormatSystemMessage(GetLastError()).c_str();
		str = str.TrimRight(_T(".!\r\n"));
		break;
	default:
		ATLASSERT(FALSE);
		str +=BkString::Get(93);
		break;
	}
	str += _T(".");

	if(nErrorCode == PWE_NOFILEACCESS_READ_KEY)
	{
		str += _T("\r\n\r\n");
		str += _T("Make sure the media that contains the key file is inserted and that KeePass has the rights to access it (check file access rights, ensure that no other application is blocking the file, ...).");
	}

	if((dwFlags & PWFF_DATALOSS_WITHOUT_SAVE) != 0)
	{
		str += _T("\r\n\r\n");
		str += _T("The target file might be in a corrupted state. Please try saving again, and if that fails, save the database to a different location.");
	}

	return str;
}


BOOL CPwUtil::AttachFileAsBinaryData(__inout_ecount(1) CPasswordProfile *pEntry,
	const TCHAR *lpFile)
{
	FILE *fp = NULL;
	LPTSTR pBinaryDesc;

	if(pEntry == NULL) return FALSE;
	ATLASSERT(lpFile != NULL); if(lpFile == NULL) return FALSE;

	_tfopen_s(&fp, lpFile, _T("rb"));
	if(fp == NULL) return FALSE;

	fseek(fp, 0, SEEK_END);
	const DWORD dwFileLen = static_cast<DWORD>(ftell(fp));
	fseek(fp, 0, SEEK_SET);

	if(dwFileLen == 0) { fclose(fp); fp = NULL; return FALSE; }
	ATLASSERT(dwFileLen > 0);

	//SAFE_DELETE_ARRAY(pEntry->pszBinaryDesc);
	SAFE_DELETE_ARRAY(pEntry->m_pAttachContent);

	DWORD i = static_cast<DWORD>(_tcslen(lpFile)) - 1;
	while(1)
	{
		if(i == (DWORD)(-1)) break;
		if((lpFile[i] == '/') || (lpFile[i] == '\\')) break;
		--i;
	}
	pBinaryDesc = (LPTSTR)&lpFile[i + 1];

	//const DWORD dwPathLen = static_cast<DWORD>(_tcslen(pBinaryDesc));

	//pEntry->pszBinaryDesc = new TCHAR[dwPathLen + 1];
	//_tcscpy_s(pEntry->pszBinaryDesc, dwPathLen + 1, pBinaryDesc);
	pEntry->m_strAttachName = CString(pBinaryDesc);

	pEntry->m_pAttachContent = new BYTE[dwFileLen];
	fread(pEntry->m_pAttachContent, 1, dwFileLen, fp);

	pEntry->m_nACLen = dwFileLen;

	fclose(fp); fp = NULL;
	return TRUE;
}

BOOL CPwUtil::SaveBinaryData(__in_ecount(1) const CPasswordProfile *pEntry,
	const TCHAR *lpFile)
{
	FILE *fp = NULL;

	if(pEntry == NULL) return FALSE;
	ATLASSERT(lpFile != NULL); if(lpFile == NULL) return FALSE;
	if(pEntry->m_strAttachName.GetLength() == 0) return FALSE;

	_tfopen_s(&fp, lpFile, _T("wb"));
	if(fp == NULL) return FALSE;

	if(pEntry->m_nACLen != 0)
		fwrite(pEntry->m_pAttachContent, 1, pEntry->m_nACLen, fp);

	fclose(fp); fp = NULL;
	return TRUE;
}

BOOL CPwUtil::RemoveBinaryData(__inout_ecount(1) CPasswordProfile *pEntry)
{
	if(pEntry == NULL) return FALSE;

	SAFE_DELETE_ARRAY(pEntry->m_pAttachContent);
	//SAFE_DELETE_ARRAY(pEntry->pszBinaryDesc);
	pEntry->m_strAttachName = L"";
	pEntry->m_nACLen = 0;
	return TRUE;
}

BOOL CPwUtil::IsAllowedStoreGroup(LPCTSTR lpGroupName, LPCTSTR lpSearchGroupName)
{
	ATLASSERT(lpGroupName != NULL); if(lpGroupName == NULL) return FALSE;

	if(_tcscmp(lpGroupName, lpSearchGroupName) == 0) return FALSE;
	return TRUE;
}

BOOL CPwUtil::IsZeroUUID(__in_ecount(16) const BYTE *pUUID)
{
	if(pUUID == NULL) return TRUE;
	return ((memcmp(pUUID, g_uuidZero, 16) == 0) ? TRUE : FALSE);
}

/* void CPwUtil::ProtectMemory(UINT8 *pMem, size_t uBytes, bool bEncrypt)
{
#ifdef _WIN32
	DWORD dwBytes = static_cast<DWORD>(uBytes);
	dwBytes += CRYPTPROTECTMEMORY_BLOCK_SIZE - (dwBytes % CRYPTPROTECTMEMORY_BLOCK_SIZE);
	ATLASSERT((uBytes <= dwBytes) && (dwBytes <= uBytes + CRYPTPROTECTMEMORY_BLOCK_SIZE));
	ATLASSERT((dwBytes % CRYPTPROTECTMEMORY_BLOCK_SIZE) == 0);

	if(bEncrypt)
	{
		ATLVERIFY(CryptProtectMemory(pMem, dwBytes, CRYPTPROTECTMEMORY_SAME_PROCESS));
	}
	else
	{
		ATLVERIFY(CryptUnprotectMemory(pMem, dwBytes, CRYPTPROTECTMEMORY_SAME_PROCESS));
	}
#endif
} */

/* CPwErrorInfo::CPwErrorInfo()
{
	m_nPwErrorCode = PWE_SUCCESS;
	m_dwLastError = 0;
}

void CPwErrorInfo::SetInfo(const std::basic_string<TCHAR>& strText,
	int nPwErrorCode, BOOL bUseLastError)
{
	m_strText = strText;
	m_nPwErrorCode = nPwErrorCode;

	if(bUseLastError == TRUE) m_dwLastError = GetLastError();
	else m_dwLastError = 0;

	this->FormatInfo();
}

void CPwErrorInfo::FormatInfo()
{
	std::basic_string<TCHAR> str;

	if(m_strText.size() > 0)
	{
		str += m_strText;
		str += _T("\r\n\r\n");
	}

	str += PWM_FormatStaticError(m_nPwErrorCode, 0);
	str += _T("\r\n\r\n");

	if(m_dwLastError != 0)
	{
		LPTSTR lpBuffer = NULL;

		if(FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
			NULL, m_dwLastError, 0, (LPTSTR)&lpBuffer, 1, NULL) != 0)
		{
			str += lpBuffer;
		}

		if(lpBuffer != NULL) { LocalFree(lpBuffer); lpBuffer = NULL; }
	}

	m_strFinal = str;
}

std::basic_string<TCHAR> CPwErrorInfo::ToString() const
{
	return m_strFinal;
}
*/

CString CPwUtil::CreateUUIDStr(CNewRandom* pRandomSource)
{
	if(pRandomSource == NULL) { ATLASSERT(FALSE); return CString(); }

	BYTE pbUuid[16];
	randCreateUUID(&pbUuid[0], pRandomSource);

	CString str;
	_UuidToString(&pbUuid[0], &str);
	return str;
}

DWORD CPwUtil::GetUniCPT(LPCTSTR lp)
{
	DWORD dwRel = 0; // See CPT specs
	while(*lp != 0) { dwRel = (ROTL32UE(dwRel, 5) + *lp); ++lp; }
	return dwRel;
}

std::basic_string<TCHAR> CPwUtil::FormatSystemMessage(DWORD dwLastErrorCode)
{
	std::basic_string<TCHAR> str;
	LPTSTR lpBuffer = NULL;

	if(FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
		NULL, dwLastErrorCode, 0, (LPTSTR)&lpBuffer, 1, NULL) != 0)
	{
		str = lpBuffer;
	}

	if(lpBuffer != NULL) { LocalFree(lpBuffer); lpBuffer = NULL; }

	return str;
}

bool CPwUtil::UnhideFile(LPCTSTR lpFile)
{
	if(lpFile == NULL) { ATLASSERT(FALSE); return false; }

	const DWORD dwAttrib = GetFileAttributes(lpFile);
	if(dwAttrib == INVALID_FILE_ATTRIBUTES) return false;

	if((dwAttrib & FILE_ATTRIBUTE_HIDDEN) == 0) return false;

	return CPwUtil::HideFile(lpFile, false);
}

bool CPwUtil::HideFile(LPCTSTR lpFile, bool bHide)
{
	if(lpFile == NULL) { ATLASSERT(FALSE); return false; }

	DWORD dwAttrib = GetFileAttributes(lpFile);
	if(dwAttrib == INVALID_FILE_ATTRIBUTES) return false;

	if(bHide) dwAttrib = ((dwAttrib & ~FILE_ATTRIBUTE_NORMAL) | FILE_ATTRIBUTE_HIDDEN);
	else // Unhide
	{
		dwAttrib &= ~FILE_ATTRIBUTE_HIDDEN;
		if(dwAttrib == 0) dwAttrib |= FILE_ATTRIBUTE_NORMAL;
	}

	return (SetFileAttributes(lpFile, dwAttrib) != FALSE);
}

CNullableEx<FILETIME> CPwUtil::GetFileCreationTime(LPCTSTR lpFile)
{
	if(lpFile == NULL) { ATLASSERT(FALSE); return CNullableEx<FILETIME>(); }

	HANDLE h = CreateFile(lpFile, GENERIC_READ, FILE_SHARE_READ, NULL,
		OPEN_EXISTING, 0, NULL);
	if(h == INVALID_HANDLE_VALUE) return CNullableEx<FILETIME>();

	FILETIME vTimes[3];
	const BOOL bResult = GetFileTime(h, &vTimes[0], &vTimes[1], &vTimes[2]);
	ATLVERIFY(CloseHandle(h));

	if(bResult == FALSE) { ATLASSERT(FALSE); return CNullableEx<FILETIME>(); }

	return CNullableEx<FILETIME>(vTimes[0]);
}

bool CPwUtil::SetFileCreationTime(LPCTSTR lpFile, const FILETIME* pTime)
{
	if(lpFile == NULL) { ATLASSERT(FALSE); return false; }

	HANDLE h = CreateFile(lpFile, FILE_WRITE_ATTRIBUTES, 0, NULL,
		OPEN_EXISTING, 0, NULL);
	if(h == INVALID_HANDLE_VALUE) return false;

	ATLVERIFY(SetFileTime(h, pTime, NULL, NULL));
	ATLVERIFY(CloseHandle(h));
	return true;
}

bool CPwUtil::EfsEncryptFile(LPCTSTR lpFile)
{
	if(lpFile == NULL) { ATLASSERT(FALSE); return false; }

	HMODULE hLib = ::LoadLibrary(_T("AdvApi32.dll"));
	if(hLib == NULL) { ATLASSERT(FALSE); return false; }

	bool bResult = false;
	LPENCRYPTFILE lpEncryptFile = (LPENCRYPTFILE)::GetProcAddress(hLib, ENCRYPTFILE_FNNAME);
	if(lpEncryptFile != NULL) { bResult = (lpEncryptFile(lpFile) != FALSE); }
	else { ATLASSERT(FALSE); }

	::FreeLibrary(hLib);
	return bResult;
}
